# -*- coding: utf-8 -*-

import tensorflow.examples.tutorials.mnist.input_data as input_data
import numpy as np
import matplotlib.pyplot as plt
import pylab
from tensorflow.python.platform import gfile
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()

def loadModel(sess):
    with gfile.FastGFile('model/keras_model.pb', 'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        sess.graph.as_default()
        tf.import_graph_def(graph_def, name='') # 导入计算图

def get_fashion_mnist_labels(index):
    text_labels = ['t-shirt', 'trouser', 'pullover', 'dress', 'coat',
                   'sandal', 'shirt', 'sneaker', 'bag', 'ankle boot']
    return text_labels[index]

def printResult(tag,result):
    maxIdx = np.argmax(result)
    print('{} = {}, label = {}'.format(tag,maxIdx,get_fashion_mnist_labels(maxIdx)))
    pass


def showPltResult(test_images,predictions,test_labels,maxIndex = 25):
    class_names = ['t-shirt', 'trouser', 'pullover', 'dress', 'coat',
                   'sandal', 'shirt', 'sneaker', 'bag', 'ankle boot']
    # 保存画布的图形，宽度为 10 ， 长度为10
    plt.figure(figsize=(10,10))
     
    # 预测 maxIndex 张图像是否准确，不准确为红色。准确为黑色
    for i in range(maxIndex):
        # 创建分布 5 * 5 个图形
        plt.subplot(5, 5, i+1)
        plt.xticks([])
        plt.yticks([])
        plt.grid(False)
        # 显示照片，以cm 为单位。
        plt.imshow(test_images[i], cmap=plt.cm.binary)
        
        # 预测的图片是否正确，黑色底表示预测正确，红色底表示预测失败
        predicted_label = np.argmax(predictions[i])
        true_label = np.argmax(test_labels[i])
        if predicted_label == true_label:
            color = 'black'
        else:
            color = 'red'
        plt.xlabel("{} ({})".format(class_names[predicted_label],
                                    class_names[true_label]),
                                    color=color)
    plt.show()
    pass

#载入数据
fashion_mnist = input_data.read_data_sets("data/", one_hot=True)
x_train, y_train = fashion_mnist.test.images, fashion_mnist.test.labels   

print('x_train.shape = ',x_train.shape);
#print('[0]=',batch_xs[0])

#使用模型判断结果
tf.reset_default_graph()

init = tf.global_variables_initializer()
sess = tf.Session()
sess.run(init)

#载入训练好的模型
loadModel(sess)

#打印网络结构
graph = tf.get_default_graph()
graph_str = "{}".format(graph.as_graph_def())
print(graph_str[0:450])

#准备要验证的图片数量
maxIndex = 25
    
#用于显示的图像数据
test_images_2_show = []
#用于验证数据结果
predictions = []
#用于对比期望的结果
test_labels = y_train[0:maxIndex]
#用于图像数据的输入
test_images2 = x_train.reshape(-1, 28, 28,1).astype('float32')  
test_images2 = test_images2[0:maxIndex]
test_images2 /= 255 # 训练时也是加入了 / 255 操作, 相当于做了一个归一化处理
#获取收集每张图片，转为 28 x 28 放入 test_images_2_show
for test_index in range(maxIndex):
    #显示图片
    one_pic_arr = np.reshape(x_train[test_index],(28,28))
    test_images_2_show.append(one_pic_arr)
    
#将数量放入模型并取得预测结果    
predictions = sess.run(graph.get_tensor_by_name("Identity:0"), feed_dict={graph.get_tensor_by_name("x:0"): test_images2}) 

#显示预测结果 
showPltResult(test_images = test_images_2_show,predictions = predictions,test_labels = test_labels,maxIndex = maxIndex)

